# -*- encoding: utf-8 -*-
"""
处理BC文件中子函数的调用问题

"""
from llvmlite import binding
import parse
import re

# 子函数调用方式目前发现3种,第三种暂无法识别
#  1.  %37 = call i32 bitcast (i32 ()* @CPBSP_FT_TestInnerIntf to i32 (i32, i32, i32*, i8*)*)(i32 %34, i32 %35, i32* %8, i8* %36), !dbg !36657
#  2.  %14 = call i32 @BDTM_GetTxPhyChNum(), !dbg !6542
#  3.  %12 = call zeroext i8 bitcast (i32 ()* @VOS_RegisterSelfTask to i8 (i32, void (i32, i32, i32, i32)*, i32)*)(i32 %11, void (i32, i32, i32, i32)* @bdtmUlLogFileSelfTask, i32 2), !dbg !42259
#      %13 = zext i8 %12 to i32, !dbg !42259
# sub_function_call_pattern = parse.compile("{space:s}{result_name:S} = call {result:S} @{function_name:w}(")
sub_function_call_pattern = parse.compile("{space:s}{result_name:S} = call {result} @{function_name:w}")
condition_pattern = parse.compile("{space:s}br {:S} {condition_name:S}, label {:S}, label {:S}")


def replace_call_sub_function_with_mock(bc_file_name, function_use_sub_functions):
    """
    把BC文件中，对子函数的调用，修改为对mock函数的调用
    1、打开bc文件
    2、解析bc文件为文本文件
    3、根据函数名称获取对应的函数代码
    4、使用mock的子函数名称代替真实的子函数名称
    5、使用修改后的函数代码，替换bc文件中原来的函数代码

    """
    with open(bc_file_name, "rb") as bc_target:
        parse_result = binding.parse_bitcode(bc_target.read())
    bc_file_text = str(parse_result)
    # print("function_use_sub_functions: %s", str(function_use_sub_functions))
    # print("replace_call_sub_function_with_mock begin")
    # print("before repalce:", bc_file_text)

    # 所有文件的函数
    for file_sub_functions in function_use_sub_functions.values():
        for function_name, sub_function_infos in file_sub_functions.items():
            # print("function_name: %s", function_name)
            original_function_text = str(parse_result.get_function(function_name))
            # print("original_function_text: ", original_function_text)
            after_replace_function_text, no_need_mock_function_list = replace_sub_function_name(original_function_text, sub_function_infos)
            bc_file_text = bc_file_text.replace(original_function_text, after_replace_function_text)
            # 非分支子函数mock，会影响ktest文件中的符号化的全局mock值，导致转dt用例时非分支子函数也会AMOCK，所以在测试入口中删除符号化的非分支子函数
            if len(no_need_mock_function_list):
                test_klee_funcname = '{}{}'.format("test_klee_", function_name)
                # 拼接测试入口定义格式 @test_klee_func()，防止有前缀字符
                if '@{}()'.format(test_klee_funcname) in bc_file_text: # 有测试入口
                    original_test_klee_function_text = str(parse_result.get_function(test_klee_funcname))
                    after_original_test_klee_function_text = original_test_klee_function_text
                    for sub_fun in no_need_mock_function_list:
                        after_original_test_klee_function_text = re.sub(r'\s+call void @klee_.+ @g_mock_{} .+'.format(sub_fun), "", after_original_test_klee_function_text)
                    # print("original_test_klee_function_text:", original_test_klee_function_text)
                    bc_file_text = bc_file_text.replace(original_test_klee_function_text, after_original_test_klee_function_text)
            # print("after repalce:", bc_file_text)
    with open("./test.klee.txt", "w") as write_text:
        write_text.write(bc_file_text)
    # 将替换后的bc内容写入原来的文件
    parse_from_new_text = binding.parse_assembly(bc_file_text)
    with open(bc_file_name, "wb") as new_bc_target:
        new_bc_target.write(parse_from_new_text.as_bitcode())
    # print("replace_call_sub_function_with_mock end")

# store load对匹配
def get_store_load_key(key:str, origin_function_text:str) -> str:
    
    condition_pattern_store_load = parse.compile("{space:s} store {:S} " + key + ",{temp}{space:s} {valueKey:S} = load {:S}, {:S}")
    condition_pattern_load_result = list(condition_pattern_store_load.findall(origin_function_text))
    # print("condition_pattern_load_result:", str(condition_pattern_load_result))
    for load_item in condition_pattern_load_result:
        # 紧挨着的store load对应该只有1个，只考虑1个
        # print("store : %s, load : %s" %(key, load_item['valueKey']))
        return str(load_item['valueKey'])


# 某种计算后的值，取等号左边的符号
def get_calc_key(key:str, origin_function_text:str) -> str:
    calc_pattern = r"\s\s(%\d+) = .*" + key
    calc_pattern_result = re.findall(calc_pattern, origin_function_text)
    for left_item in calc_pattern_result:
        # print("calc key: %s, left: %s" %(key, left_item))
        return left_item


# IR文件中语法格式
def collect_calc_sub_function_names(origin_function_text:str, icmp_names_set:set)->set:
        # 提取出所有函数调用的子函数名称和结果的变量名称
    need_mock_sub_function = set()
    # sub_function_call_pattern = parse.compile("{space:s}{result_name:S} = call {result:S} @{function_name:w}(")
    function_call_result = sub_function_call_pattern.findall(origin_function_text)
    for sub_function_item in function_call_result:
        key = sub_function_item['result_name'] 
        sub_funcvalue = sub_function_item['function_name']

        # print("key %s, sub_funcvalue: %s" %(key, sub_funcvalue))
        # 1. 0级传递，直接在分支表达式中调用函数 if (calc_a(a) < b) 
        if key in icmp_names_set:
            need_mock_sub_function.add(sub_funcvalue)
            # print("a")
            continue
        # 2. 1级传递，函数直接参与某种运算的
        calc_key_1 = get_calc_key(key, origin_function_text)
        if calc_key_1:
            # if (calc_d(b) - d)  if (c | calc_e(b))  
            if calc_key_1 in icmp_names_set:
                # print("d e")
                need_mock_sub_function.add(sub_funcvalue)
                continue
            # 2级传递，函数直接参与某种运算的返回值作为分支表达式条件
            # c += calc_c(b); if (c > 0)
            calc_load_key_1 = get_store_load_key(calc_key_1, origin_function_text)
            if calc_load_key_1 and calc_load_key_1 in icmp_names_set:
                # print("c")
                need_mock_sub_function.add(sub_funcvalue)
                continue
        # 函数返回值保存成局部变量情况，有可能局部变量多次使用
        load_key_1 = get_store_load_key(key, origin_function_text)
        if not load_key_1:
            continue
        # 3. 1级传递，函数返回值作为分支表达式  c = calc_b(a);  if (c < b)
        if load_key_1 in icmp_names_set:
            # print("b")
            need_mock_sub_function.add(sub_funcvalue)
            continue
        # 4. 2级传递，函数返回值直接在分支表达式中参与某种运算的  d = calc_f(b); if (d&c) {
        calc_key_2 = get_calc_key(load_key_1, origin_function_text)
        if not calc_key_2:
            continue
        if calc_key_2 in icmp_names_set:
            # print("f")
            need_mock_sub_function.add(sub_funcvalue)
            continue
        # 5. 3级传递，函数返回值参与某种计算后计算值保存成局部变量又参与分支表达式，d = calc_g(b); e = d / c; if ( e < 0)
        calc_load_key_3 = get_store_load_key(calc_key_2, origin_function_text)
        if calc_load_key_3 and calc_load_key_3 in icmp_names_set:
            # print("g")
            need_mock_sub_function.add(sub_funcvalue)
    return need_mock_sub_function

def collect_conditional_sub_function_names(origin_function_text):
    """
    从函数的IR代码中，提取出涉及分支判断的子函数名称
    """
    # 提取出所有条件跳转涉及的变量名称
    condition_names = set()
    branch_result = condition_pattern.findall(origin_function_text)
    for branch_item in branch_result:
        condition_names.add(branch_item['condition_name'])
    print("condition_names: ", str(condition_names))
    # 判断出涉及分支判断的函数调用，并将函数名称加入结果
    icmp_names_set = set()

    for condition_name in condition_names:
        parse_pattern = condition_name + " = icmp {:S} {:S} {value1:S}, {value2:S},"
        compare_result = list(parse.findall(parse_pattern, origin_function_text))
        if compare_result:
            icmp_names_set.add(compare_result[0]['value1'])
            icmp_names_set.add(compare_result[0]['value2'])

    need_mock_sub_function = collect_calc_sub_function_names(origin_function_text, icmp_names_set)

    print("need_mock_sub_function ：", str(need_mock_sub_function))
    # print("collect_conditional_sub_function_names end")
    return need_mock_sub_function


def replace_sub_function_name(origin_function_text, sub_function_infos):
    """
    在函数的IR代码中，把子函数的调用替换为对mock子函数的调用
    TODO 关键点：现在只替换涉及分支的子函数调用，而对一般的子函数调用，不再进行替换
    znn: 第二个参数不传递也可以
    """
    after_replace_function_text = origin_function_text
    need_replace_names_set = collect_conditional_sub_function_names(origin_function_text)
    sub_functions_set = set()
    for sub_function in sub_function_infos:
        sub_function_name = sub_function['name']
        sub_functions_set.add(sub_function_name)
        if sub_function_name in need_replace_names:
            # 2021-04-16 COMM_IniGetStrList 中调用 分支子函数COMM_IniGetStr，导致原函数也被替换成打桩，增加 ( 表调用
            mock_sub_function_name = "mock_{}(".format(sub_function_name)
            # # 防止多次替换
            # after_replace_function_text = after_replace_function_text.replace("@" + sub_function_name + "(",
            #                                                                   "@" + mock_sub_function_name)
            # # 2021-06-09 目前发现3种函数调用方式,带类型转换,替换成mock函数后报类型不匹配，所以需要把类型转换去掉
            # 1. 第一种函数调用方式带类型转换的  %53 = call i32 bitcast (i32 ()* @CPBSP_FT_CaliSendDataFromFile to i32 (i32, i8*)*)(i32 %51, i8* %52), !dbg !6490
            # 转换成mock函数后为 %53 = call i32 @mock_CPBSP_FT_CaliSendDataFromFile(i32 %51, i8* %52), !dbg !6490
            # 2. 第二种函数调用方式  %12 = call i32 @BSP_FT_GetBranchNum(i32 0, i32* %9), !dbg !6433
            # mock 函数的调用方式不带类型转换，否则写入bc文件报类型不匹配，故把待类型转换的去掉
            # 3.
            after_replace_function_text = re.sub(r'(bitcast .*?@{} to .*?\(.*?\)\*\))\(|@{}\('.format(sub_function_name, sub_function_name), "@mock_{}(".format(sub_function_name), after_replace_function_text)
    # 筛选非分支子函数，为下步 test_klee入口函数中将不是分支子函数的mock函数删除
    no_need_mock_function_list = list(filter(lambda fun: not fun.islower(), list(sub_functions_set.difference(need_replace_names_set))))

    # print("replace_sub_function_name end：",  after_replace_function_text)
    return after_replace_function_text, no_need_mock_function_list


if __name__ == '__main__':
    test_data = """
; Function Attrs: noinline nounwind
define dso_local i32 @bdtmUlLogFileSelfTaskInit() #0 !dbg !45687 {
  %1 = alloca i32, align 4
  %2 = alloca i32, align 4
  call void @llvm.dbg.declare(metadata i32* %2, metadata !45688, metadata !DIExpression()), !dbg !45689
  store i32 0, i32* %2, align 4, !dbg !45689
  call void (i8*, ...) bitcast (i32 ()* @ROSA_Printf to void (i8*, ...)*)(i8* getelementptr inbounds ([45 x i8], [45 x i8]* @.str.47.598, i32 0, i32 0)), !dbg !45690
  %3 = call i32 bitcast (i32 ()* @VOS_QueueCreate to i32 (i8*, i32, i32*, i32, i32)*)(i8* getelementptr inbounds ([15 x i8], [15 x i8]* @.str.48.599, i32 0, i32 0), i32 80, i32* @m_bdtm_ulBdtmUploadMsgCmdQID, i32 0, i32 508), !dbg !45691
  store i32 %3, i32* %2, align 4, !dbg !45692
  %4 = load i32, i32* %2, align 4, !dbg !45693
  %5 = icmp ne i32 0, %4, !dbg !45695
  br i1 %5, label %6, label %10, !dbg !45696

; <label>:6:                                      ; preds = %0
  %7 = load i32, i32* %2, align 4, !dbg !45697
  %8 = call i32 (i8*, ...) bitcast (i32 ()* @VOS_Printf to i32 (i8*, ...)*)(i8* getelementptr inbounds ([68 x i8], [68 x i8]* @.str.49.600, i32 0, i32 0), i32 %7), !dbg !45699
  %9 = load i32, i32* %2, align 4, !dbg !45700
  store i32 %9, i32* %1, align 4, !dbg !45701
  br label %21, !dbg !45701

; <label>:10:                                     ; preds = %0
  %11 = load i32, i32* @g_ulRuBdtmFid, align 4, !dbg !45702
  %12 = call zeroext i8 bitcast (i32 ()* @VOS_RegisterSelfTask to i8 (i32, void (i32, i32, i32, i32)*, i32)*)(i32 %11, void (i32, i32, i32, i32)* @bdtmUlLogFileSelfTask, i32 2), !dbg !45703
  %13 = zext i8 %12 to i32, !dbg !45703
  store i32 %13, i32* %2, align 4, !dbg !45704
  %14 = load i32, i32* %2, align 4, !dbg !45705
  %15 = icmp eq i32 255, %14, !dbg !45707
  br i1 %15, label %16, label %20, !dbg !45708

; <label>:16:                                     ; preds = %10
  %17 = load i32, i32* %2, align 4, !dbg !45709
  %18 = call i32 (i8*, ...) bitcast (i32 ()* @VOS_Printf to i32 (i8*, ...)*)(i8* getelementptr inbounds ([73 x i8], [73 x i8]* @.str.50.601, i32 0, i32 0), i32 %17), !dbg !45711
  %19 = load i32, i32* %2, align 4, !dbg !45712
  store i32 %19, i32* %1, align 4, !dbg !45713
  br label %21, !dbg !45713

; <label>:20:                                     ; preds = %10
  call void (i8*, ...) bitcast (i32 ()* @ROSA_Printf to void (i8*, ...)*)(i8* getelementptr inbounds ([43 x i8], [43 x i8]* @.str.51.602, i32 0, i32 0)), !dbg !45714
  store i32 0, i32* %1, align 4, !dbg !45715
  br label %21, !dbg !45715

; <label>:21:                                     ; preds = %20, %16, %6
  %22 = load i32, i32* %1, align 4, !dbg !45716
  ret i32 %22, !dbg !45716
}
    """
    test_data_src = """
    int calc(int a, int b)
{
    int c;
    int d;
    int e;
    if (calc_a(a) < b)
    {
        return -1;
    }
    c = calc_b(a);
    if (c < b)
    {
        return -1;
    }
    c += calc_c(b);
    
    if (c > 0)
    {
        return 0;
    }

   if (calc_d(b) - d) {
        return -1;
    }
    
    
    if (c | calc_e(b)) {
        return -1;
    }
    
    d = calc_f(b);
    
    if (d&c) {
        return -1;
    }
    
    d = calc_g(b);
    e = d / c;
    if ( e < 0) {
        return -1;
    }

    
    
    if (++b > 5) {
        return 5;
    }
    return 0;
}
"""

    test_data_1 = """
; Function Attrs: noinline nounwind
define dso_local i32 @calc(i32, i32) #0 !dbg !141 {
  %3 = alloca i32, align 4
  %4 = alloca i32, align 4
  %5 = alloca i32, align 4
  %6 = alloca i32, align 4
  %7 = alloca i32, align 4
  %8 = alloca i32, align 4
  store i32 %0, i32* %4, align 4
  call void @llvm.dbg.declare(metadata i32* %4, metadata !142, metadata !DIExpression()), !dbg !143
  store i32 %1, i32* %5, align 4
  call void @llvm.dbg.declare(metadata i32* %5, metadata !144, metadata !DIExpression()), !dbg !145
  call void @llvm.dbg.declare(metadata i32* %6, metadata !146, metadata !DIExpression()), !dbg !147
  call void @llvm.dbg.declare(metadata i32* %7, metadata !148, metadata !DIExpression()), !dbg !149
  call void @llvm.dbg.declare(metadata i32* %8, metadata !150, metadata !DIExpression()), !dbg !151
  %9 = load i32, i32* %4, align 4, !dbg !152
  %10 = call i32 @calc_a(i32 %9), !dbg !154
  %11 = load i32, i32* %5, align 4, !dbg !155
  %12 = icmp slt i32 %10, %11, !dbg !156
  br i1 %12, label %13, label %14, !dbg !157

; <label>:13:                                     ; preds = %2
  store i32 -1, i32* %3, align 4, !dbg !158
  br label %66, !dbg !158

; <label>:14:                                     ; preds = %2
  %15 = load i32, i32* %4, align 4, !dbg !160
  %16 = call i32 @calc_b(i32 %15), !dbg !161
  store i32 %16, i32* %6, align 4, !dbg !162
  %17 = load i32, i32* %6, align 4, !dbg !163
  %18 = load i32, i32* %5, align 4, !dbg !165
  %19 = icmp slt i32 %17, %18, !dbg !166
  br i1 %19, label %20, label %21, !dbg !167

; <label>:20:                                     ; preds = %14
  store i32 -1, i32* %3, align 4, !dbg !168
  br label %66, !dbg !168

返回值直接参与运算，运算结果store->load 做比较
; <label>:21:                                     ; preds = %14
  %22 = load i32, i32* %5, align 4, !dbg !170
  %23 = call i32 @calc_c(i32 %22), !dbg !171
  %24 = load i32, i32* %6, align 4, !dbg !172
  %25 = add nsw i32 %24, %23, !dbg !172
  store i32 %25, i32* %6, align 4, !dbg !172
  %26 = load i32, i32* %6, align 4, !dbg !173
  %27 = icmp sgt i32 %26, 0, !dbg !175
  br i1 %27, label %28, label %29, !dbg !176

; <label>:28:                                     ; preds = %21
  store i32 0, i32* %3, align 4, !dbg !177
  br label %66, !dbg !177
返回值直接参与运算，运算结果直接做比较
; <label>:29:                                     ; preds = %21
  %30 = load i32, i32* %5, align 4, !dbg !179
  %31 = call i32 @calc_d(i32 %30), !dbg !181
  %32 = load i32, i32* %7, align 4, !dbg !182
  %33 = add nsw i32 %31, %32, !dbg !183
  %34 = icmp ne i32 %33, 0, !dbg !183
  br i1 %34, label %35, label %36, !dbg !184

; <label>:35:                                     ; preds = %29
  store i32 -1, i32* %3, align 4, !dbg !185
  br label %66, !dbg !185
返回值直接参与运算，运算结果直接做比较
; <label>:36:                                     ; preds = %29
  %37 = load i32, i32* %6, align 4, !dbg !187
  %38 = load i32, i32* %5, align 4, !dbg !189
  %39 = call i32 @calc_e(i32 %38), !dbg !190
  %40 = or i32 %37, %39, !dbg !191
  %41 = icmp ne i32 %40, 0, !dbg !191
  br i1 %41, label %42, label %43, !dbg !192

; <label>:42:                                     ; preds = %36
  store i32 -1, i32* %3, align 4, !dbg !193
  br label %66, !dbg !193
返回值临时变量store->load到局部变量，运算结果直接做比较
; <label>:43:                                     ; preds = %36
  %44 = load i32, i32* %5, align 4, !dbg !195
  %45 = call i32 @calc_f(i32 %44), !dbg !196
  store i32 %45, i32* %7, align 4, !dbg !197
  %46 = load i32, i32* %7, align 4, !dbg !198
  %47 = load i32, i32* %6, align 4, !dbg !200
  %48 = or i32 %46, %47, !dbg !201
  %49 = icmp ne i32 %48, 0, !dbg !201
  br i1 %49, label %50, label %51, !dbg !202

; <label>:50:                                     ; preds = %43
  store i32 -1, i32* %3, align 4, !dbg !203
  br label %66, !dbg !203
返回值临时变量store->load到局部变量，运算结果store->load 做比较
; <label>:51:                                     ; preds = %43
  %52 = load i32, i32* %5, align 4, !dbg !205
  %53 = call i32 @calc_g(i32 %52), !dbg !206
  store i32 %53, i32* %7, align 4, !dbg !207
  %54 = load i32, i32* %7, align 4, !dbg !208
  %55 = load i32, i32* %6, align 4, !dbg !209
  %56 = add nsw i32 %54, %55, !dbg !210
  store i32 %56, i32* %8, align 4, !dbg !211
  %57 = load i32, i32* %8, align 4, !dbg !212
  %58 = icmp slt i32 %57, 0, !dbg !214
  br i1 %58, label %59, label %60, !dbg !215

; <label>:59:                                     ; preds = %51
  store i32 -1, i32* %3, align 4, !dbg !216
  br label %66, !dbg !216

; <label>:60:                                     ; preds = %51
  %61 = load i32, i32* %5, align 4, !dbg !218
  %62 = add nsw i32 %61, 1, !dbg !218
  store i32 %62, i32* %5, align 4, !dbg !218
  %63 = icmp sgt i32 %62, 5, !dbg !220
  br i1 %63, label %64, label %65, !dbg !221

; <label>:64:                                     ; preds = %60
  store i32 5, i32* %3, align 4, !dbg !222
  br label %66, !dbg !222

; <label>:65:                                     ; preds = %60
  store i32 0, i32* %3, align 4, !dbg !224
  br label %66, !dbg !224

; <label>:66:                                     ; preds = %65, %64, %59, %50, %42, %35, %28, %20, %13
  %67 = load i32, i32* %3, align 4, !dbg !225
  ret i32 %67, !dbg !225
}
    """
    l = collect_conditional_sub_function_names(test_data_1)
    # print(collect_conditional_sub_function_names(test_data))
    # replace_call_sub_function_with_mock("", {})
